import numpy as np
import csv
import os
from LoadTUMData import *
import sys
sys.path.append('/home/ubuntu/user_space/KF-pose/tumpreprocess')

import tumpreprocess.GtInterpolate as GI
import tumpreprocess.KFattitudeAccuracy as KA



### bianli wenjian jia neide csv wenjian, huo shoudong bianli
### xianjilu xia GT weizhi xinxi, bianlidao zhongzhi weizhi
### jisuan weizhi chazhi, jingxing jilu
### bijiao quchu zuiyouzhi

GtP = [[0.583577181, -1.177306414, 1.259987822, -1.104254476, 0.0285585595, 1.229451772],
       [-1.104254476, 0.0285585595, 1.229451772, 0.666273919, 1.107149741, 1.235215187],
       [0.666273919, 1.107149741, 1.235215187, 1.86357171, -0.479472951, 1.236903296],
       [1.86357171, -0.479472951, 1.236903296, -0.176296583, -1.182510032, 1.207645514],
       [-0.176296583, -1.182510032, 1.207645514, -0.830735894, 0.57884716, 1.556752469],
       [-0.830735894, 0.57884716, 1.556752469, 1.042996003, 1.039210659, 1.359501543],
       [1.042996003, 1.039210659, 1.359501543, 1.301728568, -0.168724015, 1.124900839],
       [1.301728568, -0.168724015, 1.124900839, -0.042418982, -1.066424815, 1.214655916],
       [-0.042418982, -1.066424815, 1.214655916, -1.132848176, 0.246177695, 1.391903119],
       [-1.132848176, 0.246177695, 1.391903119, 0.514386114, 0.86999215, 1.590409459],
       [0.514386114, 0.86999215, 1.590409459, ]]

# InitP_x = 0.583577181
# InitP_y = -1.177306414
# InitP_z = 1.259987822
# GtP_x = -1.104254476
# GtP_y = 0.0285585595
# GtP_z = 1.229451772
# Gt_xyz = np.sqrt((GtP_x - InitP_x) ** 2 + (GtP_y - InitP_y) ** 2 + (GtP_z - InitP_z) ** 2)

root = open('DistanceAnalyze.csv', 'a')
csv_writer = csv.writer(root)
csv_writer.writerow(['num', 'Gt_xyz', 'Dis_xyz',  'Raw_xyz', 'DisToGt', 'RawToGt', 'DivergeToGt', 'DivRawToGt'])

avg_DisToGt, avg_RawToGt, avg_DivergeToGt, avg_DivRawToGt = 0, 0, 0, 0
def main():
    global avg_DisToGt, avg_RawToGt, avg_DivergeToGt, avg_DivRawToGt

    time, tum_accex, tum_accey, tum_accez, tum_angux, tum_anguy, tum_anguz = loadTumImu()

    tum_px, tum_py, tum_pz, tum_q0, tum_q1, tum_q2, tum_q3 = loadTumGT()
    batchNum = 20
    for i in range(0, batchNum):
        distanceAnalyzeByGroup(time, tum_px, tum_py, tum_pz, tum_accex, tum_accey, tum_accez, tum_angux, tum_anguy, tum_anguz)
    ### for loop here

    avg_DisToGt = round(avg_DisToGt / 20, 4)
    avg_RawToGt = round(avg_RawToGt / 20, 4)
    avg_DivergeToGt = round(avg_DivergeToGt / 20, 4)
    avg_DivRawToGt = round(avg_DivRawToGt / 20, 4)
    csv_writer.writerow([avg_DisToGt, avg_RawToGt, avg_DivRawToGt, avg_DivergeToGt])


def calculateDistance():
    rootdir = '/home/ubuntu/user_space/TUMDatasetUsed'
    num = 0
    avgDisGt, avgDivGt, avgRawGt = 0, 0, 0
    for root, dirs, files in os.walk(rootdir):
        for file in files:
            PandRaw, group = fileDistance(os.path.join(root, file))
            for i in range(0, len(PandRaw)):
                row = PandRaw[i]
                # print(PandRaw[1])
                CurP_x, CurP_y, CurP_z, CurRaw_x, CurRaw_y, CurRaw_z = row[0], row[1], row[2], row[3], row[4], row[5]
                InitP_x = GtP[num][0]
                InitP_y = GtP[num][1]
                InitP_z = GtP[num][2]
                GtP_x = GtP[num][3]
                GtP_y = GtP[num][4]
                GtP_z = GtP[num][5]
                Gt_xyz = np.sqrt((GtP_x - InitP_x) ** 2 + (GtP_y - InitP_y) ** 2 + (GtP_z - InitP_z) ** 2)
                print(CurP_x, CurP_y, CurP_z)
                Dis_x = CurP_x - InitP_x
                Dis_y = CurP_y - InitP_y
                Dis_z = CurP_z - InitP_z
                Dis_xyz = np.sqrt(Dis_x * Dis_x + Dis_y * Dis_y + Dis_z * Dis_z)
                # print(Dis_xyz)
                Dis_divergeX = CurP_x - GtP_x
                Dis_divergeY = CurP_y - GtP_y
                Dis_divergeZ = CurP_z - GtP_z
                Dis_divergeXYZ = np.sqrt(
                            Dis_divergeX * Dis_divergeX + Dis_divergeY * Dis_divergeY + Dis_divergeZ * Dis_divergeZ)
                # PerDiv_Gt = Dis_xyz / Gt_xyz
                PerDis_Gt = Dis_xyz / Gt_xyz
                PerDiv_Gt = Dis_divergeXYZ / Gt_xyz

                Raw_xyz = np.sqrt((CurRaw_x - InitP_x) ** 2 + (CurRaw_y - InitP_y) ** 2 + (CurRaw_z - InitP_z) ** 2)
                PerRaw_Gt = abs((Raw_xyz - Gt_xyz) - (Dis_xyz - Gt_xyz)) / Gt_xyz

                csv_writer.writerow([file, group, round(Dis_xyz, 4), round(Dis_divergeXYZ, 4), round(Gt_xyz, 4),
                                     round(PerDis_Gt, 4), round(PerDiv_Gt, 4), round(Raw_xyz, 4), round(PerRaw_Gt, 4) ])
                avgDisGt += PerDis_Gt
                avgDivGt += PerDiv_Gt
                avgRawGt += PerRaw_Gt
                num += 1
            csv_writer.writerow([avgDisGt / num, avgDivGt / num, avgRawGt / num])




def fileDistance(file):
    sample_reader = csv.reader(open(file))
    next(sample_reader)
    curP_x, curP_y, curP_z = 0, 0, 0

    PandRaw = [[]]
### zhijie cong zheli wanghou shu 6000 weizhi, quan duchulai, cunjin yige 2D shuzu
### cnt = 6000

    cnt = 0
    for row in sample_reader:
        if ((cnt + 1) % 400 == 0) and (cnt != 0):
            print(cnt)
            curP_x, curP_y, curP_z = float(row[13]), float(row[14]), float(row[15])
            curRaw_x, curRaw_y, curRaw_z = float(row[16]), float(row[17]), float(row[18])
            PandRaw.append([curP_x, curP_y, curP_z, curRaw_x, curRaw_y, curRaw_z])

        # cnt += 1
        if cnt == 4000:
            break
        cnt += 1
    PandRaw.remove([])
    print(PandRaw)

    return PandRaw, cnt/400

def distanceAnalyzeByGroup(time, tum_px, tum_py, tum_pz, tum_accex, tum_accey, tum_accez, tum_angux, tum_anguy, tum_anguz):
    # pass in here
    # time, tum_accex, tum_accey, tum_accez, tum_angux, tum_anguy, tum_anguz = loadTumImu()
    # tum_px, tum_py, tum_pz, tum_q0, tum_q1, tum_q2, tum_q3 = loadTumGT()

    getGroup = GI.getRandomGroup(time, tum_px, tum_py, tum_pz, tum_accex, tum_accey, tum_accez, tum_angux, tum_anguy, tum_anguz)

    KFres = KA.KFfilter(getGroup)


    singleAnalyze(KFres)


def singleAnalyze(KFres):
    global avg_DisToGt, avg_RawToGt, avg_DivergeToGt, avg_DivRawToGt

    CurP_x, CurP_y, CurP_z, CurRaw_x, CurRaw_y, CurRaw_z = KFres[0], KFres[1], KFres[2], KFres[3], KFres[4], KFres[5]
    InitP_x, InitP_y, InitP_z, GtP_x, GtP_y, GtP_z = KFres[6], KFres[7], KFres[8], KFres[9], KFres[10], KFres[11]

    Gt_xyz = np.sqrt((GtP_x - InitP_x) ** 2 + (GtP_y - InitP_y) ** 2 + (GtP_z - InitP_z) ** 2)

    Dis_xyz = np.sqrt((CurP_x - InitP_x) ** 2 + (CurP_y - InitP_y) ** 2 + (CurP_z - InitP_z) ** 2)
    Raw_xyz = np.sqrt((CurRaw_x - InitP_x) ** 2 + (CurRaw_y - InitP_y) ** 2 + (CurRaw_z - InitP_z) ** 2)

    Diverge_xyz = np.sqrt((CurP_x - GtP_x) ** 2 + (CurP_y - GtP_y) ** 2 + (CurP_z - GtP_z) ** 2)
    DivRaw_xyz = np.sqrt((CurRaw_x - GtP_x) ** 2 + (CurRaw_y - GtP_y) ** 2 + (CurRaw_z - GtP_z) ** 2)

    DivergeToGt = round(Diverge_xyz / Gt_xyz, 5)
    DivRawToGt = round(DivRaw_xyz / Gt_xyz, 5)
    DisToGt = round(abs(Dis_xyz / Gt_xyz - 1), 5)
    RawToGt = round(abs(Raw_xyz / Gt_xyz - 1), 5)

    avg_DisToGt += DisToGt
    avg_RawToGt += RawToGt
    avg_DivergeToGt += DivergeToGt
    avg_DivRawToGt += DivRawToGt

    csv_writer.writerow([Gt_xyz, Dis_xyz, Raw_xyz, DisToGt, RawToGt, DivergeToGt, DivRawToGt])



if __name__ == '__main__':
    main()